import traceback
from threading import Lock, Thread
from time import sleep, time
from typing import Optional
from warnings import warn
from weakref import ref

import matplotlib.pyplot as plt
import numpy as np
from IPython.display import display
from ipywidgets import widgets


class LinesReader:
    def __init__(self, producer, consumer, *, sampling_rate=None):
        self._producer = producer
        self._consumer = consumer

        if sampling_rate:
            self._interval = 1 / sampling_rate
        else:
            self._interval = 0

        self._lock = Lock()
        self._alive = False
        self._thread = None

    @staticmethod
    def thread_func(self_ref, lock, producer, consumer, interval):
        scheduled_time = time() + interval
        while True:
            # 判断是否存活
            ctl = self_ref()
            if ctl is None:
                break
            with lock:
                if not ctl._alive:
                    break
            del ctl

            # 处理数据
            try:
                value = producer()
            except Exception:
                warn("Producer raises an exception:\n" + traceback.format_exc())
                break

            try:
                feedback = consumer(value)
            except Exception:
                warn("Consumer raises an exception:\n" + traceback.format_exc())
                break

            # 判断是否终止
            if feedback is True:
                break

            # 休眠指定时间
            now_time = time()
            sleep_time = scheduled_time - now_time
            if sleep_time > 0:
                sleep(sleep_time)
                scheduled_time = scheduled_time + interval
            else:
                scheduled_time = now_time + interval

    @property
    def alive(self):
        with self._lock:
            return self._thread and self._thread.is_alive()

    def start(self):
        with self._lock:
            if self._thread and self._thread.is_alive():
                return
            self._alive = True
            self._thread = Thread(
                target=self.thread_func,
                args=((ref(self)), self._lock, self._producer, self._consumer, self._interval),
                daemon=True)
            self._thread.start()

    def stop(self):
        with self._lock:
            if not self._thread:
                return

            self._alive = False
            thread = self._thread
            self._thread = None
        thread.join()

    def join(self):
        with self._lock:
            if not self._alive:
                return

            thread = self._thread
        thread.join()


class LinesViewer:
    def __init__(self, n: int, *, fps=64, win_size=200, hist_bins=None, figure=None):
        if not plt.isinteractive():
            warn("LinesViewer is better used in interactive mode!")

        # data
        self._n = n
        self._x_logs = []
        self._ys_logs = []
        self._win_size = win_size
        self._hist_bins = hist_bins

        # threading
        self._interval = 1 / fps
        self._lock = Lock()
        self._dirty = False
        self._thread: Optional[Thread] = None

        # init figure
        self._figure: plt.Figure = self._init_figure(figure)
        self._figure.canvas.draw_idle()

    @property
    def n(self):
        return self._n

    def append(self, x, ys):
        assert np.ndim(x) == 0

        if np.ndim(ys) == 0:
            ys = [ys]

        assert np.shape(ys) == (self._n,)

        with self._lock:
            self._x_logs.append(x)
            self._ys_logs.append(ys)

            if self._dirty:
                return
            self._dirty = True

            if self._thread:
                return
            self._thread = Thread(target=self._thread_func, daemon=True)
            self._thread.start()

    @property
    def figure(self):
        return self._figure

    def _thread_func(self):
        while True:
            start_time = time()

            with self._lock:
                x_logs = np.asarray(self._x_logs)
                ys_logs = np.asarray(self._ys_logs)
                self._dirty = False

            if len(x_logs) > 0:
                self._update_figure(x_logs, ys_logs)

            self._figure.canvas.draw_idle()

            complete_time = time()
            sleep_time = start_time + self._interval - complete_time
            if sleep_time > 0:
                sleep(sleep_time)

            with self._lock:
                if not self._dirty:
                    self._thread = None
                    break

        self._figure.canvas.draw_idle()

    def _init_figure(self, figure: Optional[plt.Figure]):
        # create Figure
        if figure is not None:
            figure = figure
        else:
            # create figure using new_figure_manager to avoid disturbing the default figure
            new_num = max(plt.get_fignums(), default=-1) + 1
            figure_manager = plt.new_figure_manager(new_num)
            figure: plt.Figure = figure_manager.canvas.figure
            if not self._hist_bins:
                figure.set_size_inches(8, 4)
            else:
                figure.set_size_inches(9, 4)
            figure.show()

        # create Axes
        if self._hist_bins:
            self._axes, self._axes_hist = figure.subplots(1, 2,
                sharey='row', gridspec_kw=dict(width_ratios=[3, 1], wspace=0.05))
        else:
            self._axes = figure.gca()
            self._axes_hist = None

        # configure Axes
        self._axes.grid()
        if self._axes_hist:
            self._axes_hist.grid()

        # create patches
        self._lines: tuple[plt.Line2D, ...] = self._axes.plot(np.zeros([0, self._n]))
        if self._axes_hist:
            self._bars: tuple[plt.Rectangle, ...] = \
                self._axes_hist.bar(np.zeros([self._hist_bins]), np.zeros([self._hist_bins]))

        return figure

    def _update_figure(self, x_logs, ys_logs):
        # adjust y range
        min_y, max_y = np.min(ys_logs), np.max(ys_logs)
        if max_y > min_y:
            padding_y = (max_y - min_y) / 32
            self._axes.set_ylim(min_y - padding_y, max_y + padding_y)
            if self._axes_hist:
                self._axes_hist.set_ylim(min_y - padding_y, max_y + padding_y)

        # lines
        for i, line in enumerate(self._lines):
            # noinspection PyTypeChecker
            line.set_xdata(x_logs)
            # noinspection PyTypeChecker
            line.set_ydata(ys_logs[:, i])

        # adjust x range
        min_x, max_x = np.min(x_logs), np.max(x_logs)
        if max_x < min_x + self._win_size:
            max_x = min_x + self._win_size
        else:
            min_x = max_x - self._win_size
        self._axes.set_xlim(min_x, max_x)

        # histogram
        if self._axes_hist:
            heights, edges = np.histogram(ys_logs[:, 0], bins=self._hist_bins)
            for i, bar in enumerate(self._bars):
                bar.set_y(edges[i])
                bar.set_height(edges[i + 1] - edges[i])
                bar.set_width(heights[i])

            # adjust x range
            self._axes_hist.set_xlim(0, np.max(heights) * 1.1)


def scan_lines_widgets(
        n, producer, *,
        sampling_rate=None,
        fps=64,
        win_size=200,
        hist_bins=None,
        figure=None,
        start=True
):
    button = widgets.Button()
    display(button)

    viewer = LinesViewer(
        n=n,
        fps=fps,
        win_size=win_size,
        hist_bins=hist_bins,
        figure=figure,
    )
    reader = LinesReader(
        producer=producer,
        consumer=lambda args: viewer.append(args[0], args[1]),
        sampling_rate=sampling_rate,
    )

    def on_button_clicked(_):
        button.disabled = True
        if reader.alive:
            reader.stop()
            button.description = 'resume'
        else:
            reader.start()
            button.description = 'stop'
        button.disabled = False

    button.description = 'start'
    button.on_click(on_button_clicked)

    if start:
        button.click()

    return reader, viewer, button
